package com.cfp4cloud.cfp.knowledge.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.Pair;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.core.util.ZipUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.http.HttpUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cfp4cloud.cfp.admin.api.feign.RemoteFileService;
import com.cfp4cloud.cfp.admin.api.utils.ByteArrayMultipartFile;
import com.cfp4cloud.cfp.knowledge.config.properties.AiKnowledgeProperties;
import com.cfp4cloud.cfp.knowledge.dto.MineruBatchResponseDTO;
import com.cfp4cloud.cfp.knowledge.dto.UmiOcrPDFModelDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiDatasetEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiDocumentEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiSliceEntity;
import com.cfp4cloud.cfp.knowledge.service.*;
import com.cfp4cloud.cfp.knowledge.support.constant.FileParserStatusEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.SliceStatusEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.SummaryStatusEnums;
import com.cfp4cloud.cfp.knowledge.support.handler.parse.UploadFileParseHandler;
import com.cfp4cloud.cfp.knowledge.support.handler.source.UploadSourceTypeHandler;
import com.cfp4cloud.cfp.knowledge.support.util.MarkdownUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.File;
import java.util.List;
import java.util.Objects;
import java.util.zip.ZipFile;

/**
 * AI定时任务服务实现类
 * <p>
 * 处理文档解析、总结和切片的异步任务 支持OCR和Mineru解析结果的更新 提供失败任务的重试机制
 *
 * @author chenda
 * @date 2024/6/16
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AiTaskServiceImpl implements AiTaskService {

	/**
	 * 文档服务
	 */
	private final AiDocumentService documentService;

	/**
	 * 数据集服务
	 */
	private final AiDatasetService datasetService;

	/**
	 * 切片服务
	 */
	private final AiSliceService sliceService;

	/**
	 * OCR助手服务
	 */
	private final AiOCRAssistantService ocrAssistantService;

	/**
	 * PDF OCR文件解析处理器
	 */
	private final UploadFileParseHandler pdfOcrUploadFileParseHandler;

	/**
	 * 上传源类型处理器
	 */
	private final UploadSourceTypeHandler uploadSourceTypeHandler;

	/**
	 * Mineru助手服务
	 */
	private final AiMineruAssistantService mineruAssistantService;

	/**
	 * 远程文件服务
	 */
	private final RemoteFileService remoteFileService;

	/**
	 * 更新文档OCR结果
	 * <p>
	 * 查询OCR解析中的文档 检查解析状态并下载成功结果 将解析内容保存到知识库
	 */
	@Override
	public void updateDocumentOcrResult() {
		List<AiDocumentEntity> documentEntityList = documentService.list(Wrappers.<AiDocumentEntity>lambdaQuery()
			.eq(AiDocumentEntity::getSliceStatus, SliceStatusEnums.OCR_PARSING.getStatus())
			.isNotNull(AiDocumentEntity::getSliceFailReason));

		for (AiDocumentEntity documentEntity : documentEntityList) {
			UmiOcrPDFModelDTO.UmiOcrPDFModelResultRequest umiOcrPDFModelResultRequest = new UmiOcrPDFModelDTO.UmiOcrPDFModelResultRequest();
			umiOcrPDFModelResultRequest.setId(documentEntity.getSliceFailReason());
			UmiOcrPDFModelDTO.UmiOcrPDFModelResultResponse umiOcrPDFModelResultResponse = ocrAssistantService
				.docResult(umiOcrPDFModelResultRequest);

			// 如果后台是成功的，则下载
			if ("success".equals(umiOcrPDFModelResultResponse.getState())) {
				UmiOcrPDFModelDTO.UmiOcrPDFModelDownRequest umiOcrPDFModelDownRequest = new UmiOcrPDFModelDTO.UmiOcrPDFModelDownRequest();
				umiOcrPDFModelDownRequest.setId(documentEntity.getSliceFailReason());
				umiOcrPDFModelDownRequest.setFile_types(List.of("txtPlain"));
				umiOcrPDFModelDownRequest.setIngore_blank(true);
				UmiOcrPDFModelDTO.UmiOcrPDFModelDownResponse umiOcrPDFModelDownResponse = ocrAssistantService
					.downloadDoc(umiOcrPDFModelDownRequest);

				TextDocumentParser textDocumentParser = new TextDocumentParser();

				// 加载文档
				Document document = UrlDocumentLoader.load(umiOcrPDFModelDownResponse.getData(), textDocumentParser);

				// 保存解析结果到知识库
				AiDatasetEntity aiDatasetEntity = datasetService.getById(documentEntity.getDatasetId());
				uploadSourceTypeHandler.saveResult(pdfOcrUploadFileParseHandler,
						Pair.of(FileParserStatusEnums.PARSE_SUCCESS, document.text()), documentEntity, aiDatasetEntity,
						null);
			}
		}

	}

	/**
	 * 更新Mineru解析结果
	 * <p>
	 * 该方法会查询所有解析中且包含失败原因的文档，并尝试重新获取解析结果。 对于解析失败的文档会更新状态为失败，对于解析成功的文档会下载并处理结果文件。
	 * </p>
	 * @throws Exception 执行过程中可能抛出的异常
	 */
	@SneakyThrows
	public void updateMineruParseResult() {
		AiKnowledgeProperties knowledgeProperties = SpringUtil.getBean(AiKnowledgeProperties.class);
		if (!knowledgeProperties.getMineru().isEnabled()) {
			return;
		}

		// 查询全部解析中的文档，并且有失败原因的 （batchId）
		List<AiDocumentEntity> parsingDocumentList = documentService.list(Wrappers.<AiDocumentEntity>lambdaQuery()
			.eq(AiDocumentEntity::getSliceStatus, SliceStatusEnums.AI_PARSING.getStatus())
			.isNotNull(AiDocumentEntity::getSliceFailReason));

		for (AiDocumentEntity documentEntity : parsingDocumentList) {
			MineruBatchResponseDTO batchResults = mineruAssistantService.getBatchResults(
					documentEntity.getSliceFailReason(),
					String.format("Bearer %s", knowledgeProperties.getMineru().getApiKey()));

			if (Objects.isNull(batchResults) || batchResults.getCode() != 0) {
				// 获取解析结果失败，更新文档状态
				documentEntity.setFileStatus(SliceStatusEnums.FAILED.getStatus());
				documentEntity.setSliceStatus(SliceStatusEnums.FAILED.getStatus());
				documentEntity.setSliceFailReason(batchResults != null ? batchResults.getMsg() : "获取解析结果失败");
				documentService.updateById(documentEntity);
				log.error("Mineru 获取解析结果失败，文档ID: {}, 错误信息: {}", documentEntity.getId(),
						batchResults != null ? batchResults.getMsg() : "未知错误");
				continue;
			}

			MineruBatchResponseDTO.ExtractResult extractResult = batchResults.getData().getExtractResult().get(0);

			if ("failed".equals(extractResult.getState())) {
				// 解析失败，更新文档状态
				documentEntity.setFileStatus(SliceStatusEnums.FAILED.getStatus());
				documentEntity.setSliceStatus(SliceStatusEnums.FAILED.getStatus());
				documentEntity.setSliceFailReason(extractResult.getErrMsg());
				documentService.updateById(documentEntity);
				log.error("Mineru 解析失败，文档ID: {}, 错误信息: {}", documentEntity.getId(), extractResult.getState());
				continue;
			}

			if ("done".equals(extractResult.getState())) {
				File tempFile = FileUtil.createTempFile();
				HttpUtil.downloadFile(extractResult.getFullZipUrl(), tempFile);
				ZipFile zipFile = new ZipFile(tempFile);

				byte[] bytes = ZipUtil.unzipFileBytes(tempFile, "full.md");
				String mdContext = StrUtil.str(bytes, CharsetUtil.defaultCharsetName());

				// 使用正则表达式匹配所有图片链接
				String result = MarkdownUtil.addDomainPrefixToImages(mdContext,
						"/api/admin/sys-file/oss/file?fileName=");

				// 处理图片资源
				for (String fileName : ZipUtil.listFileNames(zipFile, "images")) {
					byte[] imageBytes = ZipUtil.unzipFileBytes(tempFile, "images/" + fileName);
					remoteFileService.innerUpload(new ByteArrayMultipartFile(fileName, fileName, null, imageBytes),
							fileName);
				}

				FileUtil.del(tempFile);
				// 保存解析结果到知识库
				AiDatasetEntity aiDatasetEntity = datasetService.getById(documentEntity.getDatasetId());
				uploadSourceTypeHandler.saveResult(pdfOcrUploadFileParseHandler,
						Pair.of(FileParserStatusEnums.PARSE_SUCCESS, result), documentEntity, aiDatasetEntity, null);

				// 总结文档
				uploadSourceTypeHandler.summaryDocument(documentEntity, null);
				// 向量化切片
				uploadSourceTypeHandler.embedSlice(documentEntity);
			}
		}
	}

	/**
	 * 更新文档总结
	 */
	@Override
	public void updateDocumentSummary() {
		// 查询全部未总结的文档
		List<AiDocumentEntity> failSummaryDocumentList = documentService.list(Wrappers.<AiDocumentEntity>lambdaQuery()
			.eq(AiDocumentEntity::getSummaryStatus, SummaryStatusEnums.FAILED.getStatus()));

		for (AiDocumentEntity documentEntity : failSummaryDocumentList) {
			documentService.summaryDocument(documentEntity);
		}

	}

	/**
	 * 更新文档切片
	 */
	@Override
	public void updateDocumentSlice() {
		List<AiSliceEntity> aiSliceEntityList = sliceService.list(Wrappers.<AiSliceEntity>lambdaQuery()
			.eq(AiSliceEntity::getSliceStatus, SliceStatusEnums.FAILED.getStatus()));

		if (CollUtil.isEmpty(aiSliceEntityList)) {
			return;
		}

		List<Long> documentId = aiSliceEntityList.stream().map(AiSliceEntity::getDocumentId).distinct().toList();

		documentService.listByIds(documentId).forEach(documentEntity -> {
			sliceService.embedSlice(documentEntity, SliceStatusEnums.FAILED);
		});

	}

}
